#ifndef GEO_MF_LOC_BACK_CC
#define GEO_MF_LOC_BACK_CC
#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include "def.hpp"
#include "utils.hpp"
#include "topicmodel_geo_abstract.hpp"
#include "boost/unordered_map.hpp" // for hashtable

class GeoMFLocBack:public GeoTopicModel{
protected:
  double* phi_0; // background language model for the whole corpus
  vector<boost::unordered_map<int,double> > phi_geo;
  vector<boost::unordered_map<int,int> > n_r_v;
  vector<int> doc_n;
  int* geo_docs;
  double* phi_norm;  
  int* n_r;
  int* n_v;

  double* phi_norm_buffer; // used for calculating phi norm
  double lowerbound;
  double learning_rate;
  double reg;
  int N;
  ofstream file_output;
  int K; // parameter for L1 regularization by using Truncated Gradient method
  double g; // parameter for controling the sparsity
  int TEST_ERROR;
  int TEST_N;
  double* gradient_v;
public:
  GeoMFLocBack();
  virtual void Inference();
  virtual void Prediction();
  virtual double Likelihood();
  virtual int ReadDOCS();
  virtual int Interactive();
  virtual int GetDocuments(int); // this indicate how documents are read through into the model
  virtual int SaveModel();
private:
  int Single_M_Step();
  void PhiNorm();
  double ReadTest();
};

int GeoMFLocBack::SaveModel(){
  ofstream outFile;
  string  model_file = PuzaDEF::Instance()->output_file_name + ".phi0";
  outFile.open(model_file.c_str());
  for (int i=0; i<PuzaDEF::Instance()->TERM_NUM; i++){
    outFile << i << ":" << phi_0[i] << " " << endl;;
  }
  outFile.flush();
  outFile.close();

  model_file = PuzaDEF::Instance()->output_file_name + ".phi_geo";
  outFile.open(model_file.c_str());
  for(int r=0; r<PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    outFile << geo_mapping[r] << "\t";
    for(boost::unordered_map<int,double>::iterator iter = phi_geo[r].begin(); iter != phi_geo[r].end(); iter++){
      outFile << (*iter).first << ":" << (*iter).second << " ";
    }
    outFile << endl;
  }
  outFile.flush();
  outFile.close();
  return 0;
}

GeoMFLocBack::GeoMFLocBack(){
  phi_0 = new double[PuzaDEF::Instance()->TERM_NUM];
  phi_norm = new double[PuzaDEF::Instance()->GEO_REGION_NUM];
  n_r = new int[PuzaDEF::Instance()->GEO_REGION_NUM];
  phi_norm_buffer = new double[PuzaDEF::Instance()->TERM_NUM];
  n_v = new int[PuzaDEF::Instance()->TERM_NUM];
  gradient_v = new double[PuzaDEF::Instance()->TERM_NUM];
  geo_docs = new int[PuzaDEF::Instance()->GEO_REGION_NUM];

  // some initialization
  for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++){
    phi_0[i] = 0.0;
    n_v[i] = 0;
    gradient_v[i] = 0.0;
  }

  for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    n_r[r] = 0;
  }

  phi_geo.resize(PuzaDEF::Instance()->GEO_REGION_NUM);
  n_r_v.resize(PuzaDEF::Instance()->GEO_REGION_NUM);

  learning_rate = PuzaDEF::Instance()->GEO_LEARNING_RATE;
  reg = PuzaDEF::Instance()->GEO_REG;
  K = PuzaDEF::Instance()->GEO_REG_L1_K;
  g = PuzaDEF::Instance()->GEO_REG_L1_G;
}

int GeoMFLocBack::Single_M_Step(){
  while(1){
    for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
      for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
	boost::unordered_map<int,int>::const_iterator find_iter = n_r_v[r].find(v);
	double temp  = n_r[r] * exp(phi_0[v] + phi_geo[r][v] - phi_norm[r]);
	if(find_iter != n_r_v[r].end()){
	  double gradient_r_v = (*find_iter).second - temp;
	  phi_geo[r][v] = Trunc(phi_geo[r][v] + learning_rate * gradient_r_v, g);
	}
	else{
	  double gradient_r_v = 0 - temp;
	  phi_geo[r][v] = Trunc(phi_geo[r][v] + learning_rate * gradient_r_v, g);
	}
	gradient_v[v] += temp;
      }
    }
    
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      phi_0[v] = phi_0[v] + learning_rate * (n_v[v] - gradient_v[v]);
      gradient_v[v] = 0.0;
    }
    break;
  }
  return 0;
}



void GeoMFLocBack::Inference(){
  PuzaLogger::Instance()->PutString("Start Inference for Simple Geo Matrix Factorization Model [Loc + Back]");
  int ITER = 0;
  double old_lowerbound = 0;
  while(1){
    PuzaLogger::Instance()->PutString("Iteration "+to_string<int>(ITER,std::dec));
    Single_M_Step();
    Likelihood();
    PuzaLogger::Instance()->PutString("\tLikelihood:" + to_string<double>(lowerbound,std::dec));
    if((lowerbound < old_lowerbound) && ( ITER != 0))
      break;
    old_lowerbound = lowerbound;
    if (ITER>=500){
      break;
    }
    if (ITER % 50 == 0){
      SaveModel();
    }
    ITER ++;
  }
  PuzaLogger::Instance()->PutString("Batch Learning Finished.");
}

void GeoMFLocBack::Prediction(){
}
 
double GeoMFLocBack::Likelihood(){
  double result = 0.0;

  for(unsigned int d = 0; d < TM_DOCS.size(); d++){
    int region = doc_geo_stamp[d];
    TermList current_list = TM_DOCS[d];
    int term_count = current_list.size();
    for (int j=0;j < term_count; j++){
      int term_id = current_list[j].term_id;
      int count = current_list[j].term_stat;
      result = result + count * (phi_0[term_id] + phi_geo[region][term_id] - phi_norm[region]);
    }      
  }
  lowerbound = result;
  return lowerbound;
}

double GeoMFLocBack::ReadTest(){
  double likelihood = 0.0;
  string temp;
  if(PuzaDEF::Instance()->input_file_name.empty()){
    cout << "Please give input file name." << endl;
    exit(0);
  }
  ifstream testFile;
  testFile.open(PuzaDEF::Instance()->input_file_name.c_str());
  if(!testFile){
    cout << PuzaDEF::Instance()->input_file_name << endl;
    cout << "Can't open the file." << endl;
    exit(0);
  }
  string original_id;
  string temp_geo_stamp;
  string temp_user_stamp;
  vector<string> buffer;
  TEST_ERROR = 0;
  TEST_N = 0;
  while(!testFile.eof()){
    getline(testFile,temp,'\n');
    std::istringstream iss(temp);
    // get the doc id
    getline(iss,temp,'\t');
    original_id = temp;
    // get geo                                                                                                                                 
    getline(iss,temp,'\t');
    temp_geo_stamp = temp;
    // get user
    getline(iss,temp,'\t');
    temp_user_stamp = temp;
    // get the doc body
    getline(iss,temp,'\n');
    std::istringstream tempiss(temp);
    vector<string> tokens;
    copy(istream_iterator<string>(tempiss), istream_iterator<string>(), back_inserter<vector<string> >(tokens));
    if(tokens.size()<1){
      continue;
    }

    boost::unordered_map<string,int>::iterator geo_iter = geo_reverse_mapping.find(temp_geo_stamp);
    int geo_id = -1;
    if(geo_iter != geo_reverse_mapping.end()){
      // id found
      geo_id = (*geo_iter).second;
    }
    else{
      geo_id = geo_reverse_mapping.size();
      if(geo_id >= PuzaDEF::Instance()->GEO_REGION_NUM){
	cout << "region wrong" << endl;
	exit(0);
      }
      geo_reverse_mapping[temp_geo_stamp] = geo_id;
      geo_mapping[geo_id] = temp_geo_stamp;
    }
    TermList test_doc;
    for (vector<string>::iterator iter = tokens.begin(); iter < tokens.end(); iter++){
      int word_id;
      int word_count;
      std::istringstream valueiss((*iter));
      getline(valueiss,temp,':');
      from_string<int>(word_id,temp,std::dec);
      getline(valueiss,temp);  // get count
      from_string<int>(word_count,temp,std::dec);
      TermItem newAssign;
      newAssign.term_id = word_id;
      newAssign.term_stat = word_count;
      test_doc.push_back(newAssign);
    }
    // prediction
    double max_likelihood = 0.0;
    int max_r = -1;
    int term_count = test_doc.size();
    for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
      double temp_likelihood = 0.0;
      for(int j=0; j < term_count; j++){
	int term_id = test_doc[j].term_id;
	int count = test_doc[j].term_stat;
	temp_likelihood+= count * (phi_0[term_id] + phi_geo[r][term_id] - phi_norm[r]);
      }
      if(r == 0){
	max_likelihood = temp_likelihood;
      }
      else{
	if(temp_likelihood > max_likelihood){
	  max_r = r;
	  max_likelihood = temp_likelihood;
	}
      }	
    }
    // calculate likelihood
    for(int j=0; j < term_count; j++){
      int term_id = test_doc[j].term_id;
      int count = test_doc[j].term_stat;
      likelihood+= count * (phi_0[term_id] + phi_geo[geo_id][term_id] - phi_norm[geo_id]);
    }
    if(max_r != geo_id)
      TEST_ERROR++;
    TEST_N++;  
  }
  testFile.close();
  return likelihood;
}

/* We train the model and do prediction interleavely */
int GeoMFLocBack::Interactive(){
  PuzaLogger::Instance()->PutString("Start Interactive Inference for Simple Geo Matrix Factorization Model [Only Loc]");
  int ITER = 0;
  while(1){
    PuzaLogger::Instance()->PutString("Iteration "+to_string<int>(ITER,std::dec));
    Single_M_Step();
    Likelihood();
    PuzaLogger::Instance()->PutString("\tTraining Set Likelihood:" + to_string<double>(lowerbound,std::dec));
    double test_likelihood = ReadTest();
    double error_rate = (double)TEST_ERROR/TEST_N;
    PuzaLogger::Instance()->PutString("\tTest Set Accuracy:" + to_string<double>(1-error_rate,std::dec));
    PuzaLogger::Instance()->PutString("\tTest Set Likelihood:" + to_string<double>(test_likelihood,std::dec));
    if (ITER>=500){
      break;
    }
    if (ITER % 20 == 0){
      SaveModel();
    }
    ITER ++;
  }
  PuzaLogger::Instance()->PutString("Interactive Inference Finished.");
  return 0;
}

int GeoMFLocBack::ReadDOCS(){
  int MAX_INT = 0;
  N = 0;
  string temp;
  if(PuzaDEF::Instance()->input_file_name.empty()){
    cout << "Please give input file name." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->InputFile.open(PuzaDEF::Instance()->input_file_name.c_str());
  if(!PuzaDEF::Instance()->InputFile){
    cout << "Can't open the file." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->data_ids.clear();
  string original_id;
  string temp_geo_stamp;
  string temp_user_stamp;
  vector<string> buffer;
  doc_n.clear();
  while(!PuzaDEF::Instance()->InputFile.eof()){
    getline(PuzaDEF::Instance()->InputFile,temp,'\n');
    std::istringstream iss(temp);
    // get the doc id
    getline(iss,temp,'\t');
    original_id = temp;
    // get geo                                                                                                                                 
    getline(iss,temp,'\t');
    temp_geo_stamp = temp;
    // get user
    getline(iss,temp,'\t');
    temp_user_stamp = temp;
    // get the doc body
    getline(iss,temp,'\n');
    std::istringstream tempiss(temp);
    vector<string> tokens;
    copy(istream_iterator<string>(tempiss), istream_iterator<string>(), back_inserter<vector<string> >(tokens));
    if(tokens.size()<1){
      continue;
    }

    int now_id = PuzaDEF::Instance()->data_ids.size();
    PuzaDEF::Instance()->data_ids[now_id] = original_id;
    boost::unordered_map<string,int>::iterator geo_iter = geo_reverse_mapping.find(temp_geo_stamp);
    int geo_id = -1;
    if(geo_iter != geo_reverse_mapping.end()){
      // id found
      geo_id = (*geo_iter).second;
      doc_geo_stamp.push_back(geo_id);
    }
    else{
      geo_id = geo_reverse_mapping.size();
      geo_reverse_mapping[temp_geo_stamp] = geo_id;
      geo_mapping[geo_id] = temp_geo_stamp;
      doc_geo_stamp.push_back(geo_id);
    }
    TermList assignList;
    int temp_n = 0;
    for (vector<string>::iterator iter = tokens.begin(); iter < tokens.end(); iter++){
      int word_id;
      int word_count;
      std::istringstream valueiss((*iter));
      getline(valueiss,temp,':');
      from_string<int>(word_id,temp,std::dec);
      getline(valueiss,temp);  // get count
      from_string<int>(word_count,temp,std::dec);
      TermItem newAssign;
      newAssign.term_id = word_id;
      newAssign.term_stat = word_count;
      assignList.push_back(newAssign);
      n_v[word_id] = n_v[word_id] + word_count;
      n_r_v[geo_id][word_id] = n_r_v[geo_id][word_id] + word_count;
      N = N + word_count;
      temp_n = temp_n + word_count;
    }
    if((int)assignList.size() > MAX_INT)
      MAX_INT = assignList.size();
    TM_DOCS.push_back(assignList);
    doc_n.push_back(temp_n);
    n_r[geo_id] += temp_n;
  }
  PuzaDEF::Instance()->InputFile.close();
  return MAX_INT;
}

void GeoMFLocBack::PhiNorm(){
  double* buffer = new double[PuzaDEF::Instance()->TERM_NUM];
  for(int r = 0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++){
      buffer[i] = phi_0[i] + phi_geo[r][i];
    }
    phi_norm[r] = log_sum_array(buffer,PuzaDEF::Instance()->TERM_NUM); 
  }
  delete[] buffer;
}

int GeoMFLocBack::GetDocuments(int choice){
  ReadDOCS();
  for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++)
    phi_0[i] = log((n_v[i] + 1.0)/(N+PuzaDEF::Instance()->TERM_NUM));

  for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    //for(boost::unordered_map<int,int>::iterator iter = n_r_v[r].begin(); iter != n_r_v[r].end(); iter++){
    //      int term_id = (*iter).first;
    //      double value = (*iter).second;
    //      double temp = log(value / n_r[r]) - phi_0[term_id];
    //      phi_geo[r][term_id] = temp;
    //    }
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      boost::unordered_map<int,int>::const_iterator iter = n_r_v[r].find(v);
      if(iter != n_r_v[r].end()){
        phi_geo[r][v] = log(((*iter).second + 1.0)/(n_r[r] + PuzaDEF::Instance()->TERM_NUM)) - phi_0[v];
      }
      else{
        phi_geo[r][v] = log(1.0/(n_r[r] + PuzaDEF::Instance()->TERM_NUM)) - phi_0[v];
      }
    }
  }
  PhiNorm();
  return 0;
}
#endif
